package org.nicta.datasetLoader.MovieLens;

import java.io.BufferedReader;
import java.util.HashMap;

public class MovieLensHelper {
	public static final String small ="100K";
	public static final String medium ="1M";
	public static final String large ="10M";

	// datasetType {0:"100k", 1:"1M", 2:"10M"}
	public static MovieLensInfo getInfo(String datasetType){
		MovieLensInfo info = null;
		if(datasetType.equals(small)){
			info = load100KInfo();
		}
		else if (datasetType.equals(medium)){
			info = load1MInfo();
		}
		else if(datasetType.equals(large)){
			info = load10MInfo();
		}
		return info;
	}

	private static MovieLensInfo load100KInfo() {
		MovieLensInfo info = new MovieLensInfo("100K");
		
		String[] occupation = {
	            "administrator",
	            "artist",
	            "doctor",
	            "educator",
	            "engineer",
	            "entertainment",
	            "executive",
	            "healthcare",
	            "homemaker",
	            "lawyer",
	            "librarian",
	            "marketing",
	            "none",
	            "other",
	            "programmer",
	            "retired",
	            "salesman",
	            "scientist",
	            "student",
	            "technician",
	            "writer"
	    };
		 final String[] genre = {
	             "unknown",
	             "Action",
	             "Adventure",
	             "Animation",
	             "Children's",
	             "Comedy",
	             "Crime",
	             "Documentary",
	             "Drama",
	             "Fantasy",
	             "Film-Noir",
	             "Horror",
	             "Musical",
	             "Mystery",
	             "Romance",
	             "Sci-Fi",
	             "Thriller",
	             "War",
	             "Western"
	     };
		 info.occupation = occupation;
		 info.genre = genre;
		 info.usersCount = 943;
		 info.moviesCount = 1682;
		 info.separator = "\\|";
//		 info.userFeaturesCount = info.usersCount+10+1+info.occupation.length;
//		 info.movieFeaturesCount = info.moviesCount+19;
		 info.userFeaturesCount = 10+1+info.occupation.length;
		 info.movieFeaturesCount = genre.length;
//		 info.userFeaturesCount = 10+1+info.occupation.length;
//		 info.movieFeaturesCount = 19;
		 return info;
	} 

	private static MovieLensInfo load1MInfo() {
		
		MovieLensInfo info = new MovieLensInfo("1M");
		
		String[] occupation = {
				"other",
				"academic",
				"artist",
				"clerical",
				"college student",
				"customer service",
				"health care",
				"managerial",
				"farmer",
				"homemaker",
				"K-12 student",
				"lawyer",
				"programmer",
				"retired",
				"marketing",
				"scientist",
				"self-employed",
				"engineer",
				"craftsman",
				"unemployed",
				"writer"
				};
		
		 final String[] genre = {
				 "Action",
				 "Adventure",
				 "Animation",
				 "Children's",
				 "Comedy",
				 "Crime",
				 "Documentary",
				 "Drama",
				 "Fantasy",
				 "Film-Noir",
				 "Horror",
				 "Musical",
				 "Mystery",
				 "Romance",
				 "Sci-Fi",
				 "Thriller",
				 "War",
				 "Western"
				 };
		 
		 info.occupation = occupation;
		 info.genre = genre;
		 
		 info.usersCount = 6040;
		 info.moviesCount = 3952;
		 info.separator = "::";
		 info.userFeaturesCount = 7+1+info.occupation.length;
		 info.movieFeaturesCount = genre.length;
		 //total ratings = 1000209
		 return info;
		
	}

	private static MovieLensInfo load10MInfo() {
		/* TODO
		 * - return MovieLensInfo object relevant to 10M dataset
		 */
		return null;
	}
	
	private static HashMap<Integer[], Double> loadRatings(
			BufferedReader reader, MovieLensInfo info) throws Exception{
		HashMap<Integer[], Double> ratings = new HashMap<Integer[], Double>();
		String line = reader.readLine();
		while (line != null) {
			String[] tokens = line.split(info.separator);
	        int userId = Integer.parseInt(tokens[0]);
	        int movieId = Integer.parseInt(tokens[1]);
	        double rating = Float.parseFloat(tokens[2]);
	        ratings.put(new Integer[]{userId, movieId}, rating);
	         
	        line = reader.readLine();
		         
		 }
		return ratings;
	}
	
	
	public static HashMap<Integer[], Double> load100KRatings(
			BufferedReader reader, MovieLensInfo info) throws Exception{
		
		HashMap<Integer[], Double> ratings;
		ratings = loadRatings(reader, info);
		return ratings;
	}

	public static HashMap<Integer[], Double> load1MRatings(
			BufferedReader reader, MovieLensInfo info) throws Exception {
		HashMap<Integer[], Double> ratings;
		ratings = loadRatings(reader, info);
		return ratings;
	}

	public static HashMap<Integer[], Double> load10MRatings(
			BufferedReader ratingReader, MovieLensInfo info) throws Exception {
		// TODO Auto-generated method stub
		return null;
	}

	public static HashMap<Integer, double[]> load100KUserFeatures(
			BufferedReader reader, MovieLensInfo info) throws Exception{
        
		HashMap<Integer, double[]> userFeatures = 
				new HashMap<Integer, double[]>();
				
		String line = reader.readLine();
        while (line != null) {
    		double[] features = new double[info.userFeaturesCount];
            String[] tokens = line.split(info.separator);
            int decade = (int)(Double.parseDouble(tokens[1]) / 10d);
            for (int i = 0; i < 10; i++)
                    features[i] = (decade == i) ? 1d : 0d;

            features[10] = tokens[2].equals("M") ? 1d : 0d;
    
            for (int x = 0; x < info.occupation.length; x++)
                    features[x + 11] = 
                    		(info.occupation[x].equals(tokens[3])) ? 1d : 0d;
            
//            int offset = info.occupation.length + 10 + 1 - 1;
//            int userId = (int) Double.parseDouble(tokens[0]);
//            features[userId+offset-1] = 1; //userid starts with 1
            
            userFeatures.put(Integer.parseInt(tokens[0]), features);
            line = reader.readLine();
        }
        
		return userFeatures;
	}

	public static HashMap<Integer, double[]> load1MUserFeatures(
			BufferedReader reader, MovieLensInfo info)throws Exception {
		
		HashMap<Integer, double[]> userFeatures = 
				new HashMap<Integer, double[]>();

		HashMap<Integer, Integer> ageToIndex = new HashMap<Integer, Integer>();
		
		ageToIndex.put(1,0);
		ageToIndex.put(18,1);
		ageToIndex.put(25,2);
		ageToIndex.put(35,3);
		ageToIndex.put(45,4);
		ageToIndex.put(50,5);
		ageToIndex.put(56,6);

		String line = reader.readLine();
        while (line != null) {
    		double[] features = new double[info.userFeaturesCount];
            String[] tokens = line.split(info.separator);
            Integer ageGroup = Integer.parseInt(tokens[2]);
            Integer index = ageToIndex.get(ageGroup);
            features[index] = 1d;
            features[7] = tokens[1].equals("M") ? 1d : 0d;
            Integer occIndex = Integer.parseInt(tokens[3]);
            features[8+occIndex] = 1d;
            userFeatures.put(Integer.parseInt(tokens[0]), features);
            line = reader.readLine();
        }
        
		return userFeatures;
	}

	public static HashMap<Integer, double[]> load10MUserFeatures(
			BufferedReader reader, MovieLensInfo info)throws Exception {
		// TODO Auto-generated method stub
		return null;
	}

	public static HashMap<Integer, double[]> load100KMovieFeatures(
			BufferedReader reader, MovieLensInfo info)throws Exception {

		HashMap<Integer, double[]> movieFeatures = 
				new HashMap<Integer, double[]>();
				
		String line = reader.readLine();
		while(line != null){
			  double[] features = new double[info.movieFeaturesCount];
              String[] tokens = line.split(info.separator);
              for (int x = 5; x < tokens.length; x++) {
                  features[x-5] = Double.parseDouble(tokens[x]);
              }
//              int offset = tokens.length - 5 - 1;
//              int movieId = (int) Double.parseDouble(tokens[0]);
//              features[movieId+ offset - 1] = 1; //movieid Starts with 1
//              
              
              movieFeatures.put(Integer.parseInt(tokens[0]), features);
              line = reader.readLine();
		}	
		return movieFeatures;
	}

	public static HashMap<Integer, double[]> load1MMovieFeatures(
			BufferedReader reader, MovieLensInfo info)throws Exception  {
		
		HashMap<Integer, double[]> movieFeatures = 
				new HashMap<Integer, double[]>();
		HashMap<String, Integer> genreToIndex = new HashMap<String, Integer>();	
		for(int i = 0; i < info.genre.length; i++){
			genreToIndex.put(info.genre[i], i);
		}
		int index;
		String line = reader.readLine();
		while(line != null){
			  double[] features = new double[info.movieFeaturesCount];
              String[] tokens = line.split(info.separator);
              for (int x = 2; x < tokens.length; x++) {
                  index = genreToIndex.get(tokens[x]);
                  features[index] = 1d;
              }
              movieFeatures.put(Integer.parseInt(tokens[0]), features);
              line = reader.readLine();
		}
		return movieFeatures;

	}

	public static HashMap<Integer, double[]> load10MMovieFeatures(
			BufferedReader reader, MovieLensInfo info)throws Exception  {
		// TODO Auto-generated method stub
		return null;	}
}
